import codecs
import numpy as np
from gensim import corpora
from gensim.models import LdaModel
from gensim.models.coherencemodel import CoherenceModel
import matplotlib.pyplot as plt
import os
from multiprocessing import freeze_support

# 设置matplotlib中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

def compute_coherence_values(texts, dictionary, corpus, start=1, end=16, step=1):
    coherence_values = []
    for num_topics in range(start, end, step):
        model = LdaModel(corpus=corpus, 
                        id2word=dictionary,
                        num_topics=num_topics, 
                        random_state=42,
                        passes=60)
        
        coherence_model = CoherenceModel(model=model, 
                                       texts=texts,
                                       dictionary=dictionary,
                                       coherence='c_v')
        
        coherence_values.append(coherence_model.get_coherence())
        print(f'主题数 {num_topics} 的一致性得分: {coherence_values[-1]}')
    
    return coherence_values

def main():
    try:
        # 获取当前脚本所在的目录
        current_dir = os.path.dirname(os.path.abspath(__file__))
        # 获取项目根目录
        project_dir = os.path.dirname(current_dir)
        # 构建输入文件的完整路径
        input_file = os.path.join(project_dir, 'CNKI-output.txt')
        
        print(f"尝试读取文件: {input_file}")
        
        # 读取数据
        train = []
        with codecs.open(input_file, 'r', encoding='utf8') as f:
            for line in f:
                if line.strip():
                    train.append(line.strip().split())

        if not train:
            raise ValueError("没有读取到有效数据")

        # 创建词典和语料库
        dictionary = corpora.Dictionary(train)
        corpus = [dictionary.doc2bow(text) for text in train]

        # 计算不同主题数的一致性得分
        coherence_values = compute_coherence_values(train, dictionary, corpus)

        # 找出局部最大值
        local_maxima = []
        for i in range(1, len(coherence_values)-1):
            if coherence_values[i] > coherence_values[i-1] and coherence_values[i] > coherence_values[i+1]:
                local_maxima.append((i+1, coherence_values[i]))

        # 创建输出目录
        output_dir = os.path.join(current_dir)
        os.makedirs(output_dir, exist_ok=True)

        # 绘制图形
        plt.figure(figsize=(10, 6))
        plt.plot(range(1, 16), coherence_values, 'b-', marker='o')
        plt.xlabel('主题数量')
        plt.ylabel('一致性得分')
        plt.title('主题数量与一致性得分关系')
        plt.grid(True)

        # 标记局部最大值点
        for topic_num, coherence in local_maxima:
            plt.plot(topic_num, coherence, 'ro', markersize=10)
            plt.annotate(f'主题数={topic_num}\n得分={coherence:.4f}', 
                        xy=(topic_num, coherence), 
                        xytext=(10, 10),
                        textcoords='offset points')

        # 保存图形
        plot_file = os.path.join(output_dir, 'coherence_scores.png')
        plt.savefig(plot_file, dpi=300, bbox_inches='tight')
        plt.close()

        # 保存结果
        results_file = os.path.join(output_dir, 'coherence_results.txt')
        with open(results_file, 'w', encoding='utf-8') as f:
            f.write('主题一致性分析结果\n')
            f.write('=' * 50 + '\n\n')
            
            f.write('局部最大值：\n')
            for topic_num, coherence in local_maxima:
                f.write(f'主题数: {topic_num}, 一致性得分: {coherence:.4f}\n')
            
            f.write('\n所有主题数的一致性得分：\n')
            for i, score in enumerate(coherence_values, 1):
                f.write(f'主题数: {i}, 一致性得分: {score:.4f}\n')
            
            # 找出全局最大值
            best_num_topics = np.argmax(coherence_values) + 1
            f.write(f'\n最优主题数（全局最大值）: {best_num_topics}\n')
            f.write(f'对应一致性得分: {max(coherence_values):.4f}\n')

        print(f'\n分析完成！')
        print(f'结果已保存到: {results_file}')
        print(f'可视化图表已保存到: {plot_file}')

    except Exception as e:
        print(f"错误: {str(e)}")
        print(f"当前工作目录: {os.getcwd()}")

if __name__ == '__main__':
    freeze_support()  # 添加多进程支持
    main() 